package main

import (
	"fmt"
	"sync"
	"time"
)

/*
限流器：控制请求频率，保护系统
*/
type RateLimiter struct {
	tokens    chan struct{}
	rate      time.Duration
	closeChan chan struct{}
}

func NewRateLimiter(rate time.Duration, burst int) *RateLimiter {
	rl := &RateLimiter{
		tokens:    make(chan struct{}, burst),
		rate:      rate,
		closeChan: make(chan struct{}),
	}

	// 定期添加token
	go func() {
		ticker := time.NewTicker(rate)
		defer ticker.Stop()

		for {
			select {
			case <-ticker.C:
				select {
				case rl.tokens <- struct{}{}:
				default: // token池已满，丢弃
				}
			case <-rl.closeChan:
				return
			}
		}
	}()

	return rl
}

func (rl *RateLimiter) Allow() bool {
	select {
	case <-rl.tokens:
		return true
	default:
		return false
	}
}

func (rl *RateLimiter) Wait() {
	<-rl.tokens
}

func (rl *RateLimiter) Close() {
	close(rl.closeChan)
}

func main() {
	// 限制为每秒2个请求，突发最多5个
	limiter := NewRateLimiter(500*time.Millisecond, 5)
	defer limiter.Close()

	var wg sync.WaitGroup

	for i := 1; i <= 10; i++ {
		wg.Add(1)
		go func(id int) {
			defer wg.Done()

			fmt.Printf("任务 %d 等待限流器...\n", id)
			limiter.Wait() // 等待获得token

			fmt.Printf("任务 %d 开始执行 [%s]\n", id, time.Now().Format("15:04:05.000"))
			time.Sleep(100 * time.Millisecond) // 模拟任务执行
			fmt.Printf("任务 %d 完成\n", id)
		}(i)

		time.Sleep(100 * time.Millisecond)
	}

	wg.Wait()
}
